package cz.cuni.amis.episodic.bayes.memories;

import java.io.File;
import java.io.IOException;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.apache.commons.lang.StringUtils;
import org.apache.commons.math.stat.descriptive.AggregateSummaryStatistics;
import org.apache.commons.math.stat.descriptive.DescriptiveStatistics;

import com.itextpdf.text.pdf.PdfStructTreeController.returnType;

import cz.cuni.amis.episodic.bayes.memories.eval.RecallResult;
import cz.cuni.amis.episodic.bayes.memories.eval.SingleTestCasePerformance;
import cz.cuni.amis.episodic.bayes.utils.DyBaNeMUtil;
import cz.cuni.amis.episodic.bayes.utils.chart.GraphPaintingDevice;
import cz.cuni.amis.episodic.data.EpisodicHierarchyInstance;
import cz.cuni.amis.episodic.dybanem.DBN;
import cz.cuni.amis.episodic.dybanem.RetrospectiveMemoriesCreator;

/**
 * Set of methods for evaluating strategies.
 * 
 * @author ik
 */
public class MemoryCreationStrategy {

	/*
	 * static SingleTestCasePerformance computePerformance( List<List<String>>
	 * groundTruth, DBN dbn, List<MemoryElement> mems) { computePerformance(
	 * List<List<String>> groundTruth, DBN dbn, List<MemoryElement> mems, in) }
	 */

	/**
	 * Tests how well mems reconstruct the original sequence.
	 * 
	 * @param groundTruth
	 * @param dbn
	 * @param mems
	 * @return
	 */
	static SingleTestCasePerformance computePerformance(
			EpisodicHierarchyInstance groundTruth, DBN dbn,
			List<MemoryElement> mems) {

		double[] performance = new double[mems.size()];
		List<RecallResult> results = new ArrayList<>();

		for (int i = 0; i < mems.size(); i++) {
			List<MemoryElement> memories = mems.subList(0, i + 1);
			dbn.resetEvidence(memories);
			dbn.getNet().updateBeliefs();
			EpisodicHierarchyInstance recall = getRecall(dbn, groundTruth
					.getHierarchyByTime().size());
			results.add(computeRecallResult(groundTruth, recall));
		}
		return new SingleTestCasePerformance(mems, results);
	}

	public static RecallResult computeRecallResult(
			EpisodicHierarchyInstance groundTruth,
			EpisodicHierarchyInstance recall) {
		// compute performance for every level
		double[] performance = new double[groundTruth.getHierarchyByTime()
				.size()];
		for (int i = 0; i < performance.length; i++) {
			performance[i] = computeMatch(
					groundTruth.getHierarchyByTime().get(i), recall
							.getHierarchyByTime().get(i));
		}
		return new RecallResult(recall, performance);
	}

	public static EpisodicHierarchyInstance getRecall(DBN dbn, int depth) {
		List<List<String>> recalls = new ArrayList<>();
		for (int i = 0; i < depth; i++) {
			recalls.add(Arrays.asList(dbn
					.getDBNNodeMostProbableStringValues(DyBaNeMUtil
							.getNodeNameForDepth(i))));
		}

		return new EpisodicHierarchyInstance(recalls);
	}

	public static void outputResults(PrintStream out,
			List<SingleTestCasePerformance> performances) {
		for (SingleTestCasePerformance perf : performances) {
			out.println(perf.toString());
		}
	}

	/**
	 * [hierarchyLevel][numberOfMems]
	 * @param performances
	 * @return
	 */
	public static DescriptiveStatistics[][] agregatePerformances(List<SingleTestCasePerformance> performances) {
		int maxHierDepth = performances.get(0).getResults().get(0).getPerformanceByLevel().length;
		DescriptiveStatistics[][] res = new DescriptiveStatistics[maxHierDepth][];
		for(int i = 0; i < maxHierDepth; i++) {
			res[i] = agregatePerformances(performances, i);
		}
		
		return res;
	}
		
	/**
	 * Computes performance statistics for 1, 2, .. n mems.
	 * 
	 * @param performances
	 * @return
	 */
	public static DescriptiveStatistics[] agregatePerformances(
			List<SingleTestCasePerformance> performances, int hierarchyLevel) {
		int depth = performances.get(0).getMems().size();
		DescriptiveStatistics[] performance = new DescriptiveStatistics[depth];
		for (int i = 0; i < performance.length; i++) {
			performance[i] = new DescriptiveStatistics();
		}

		for (SingleTestCasePerformance perf : performances) {
			for (int i = 0; i < depth; i++) {
				performance[i].addValue(perf.getResults().get(i).getPerformanceByLevel()[hierarchyLevel]);
			}
		}

		return performance;
	}

	/**
	 * Computes computation time statistics.
	 * 
	 * @param performances
	 * @return
	 */
	public static DescriptiveStatistics agregateComputationTime(
			List<SingleTestCasePerformance> performances) {
		//int depth = performances.get(0).getPerformance().length;
		DescriptiveStatistics performance = new DescriptiveStatistics();

		for (SingleTestCasePerformance perf : performances) {
			performance.addValue(perf.getTotalComputationTimeNanos());
		}

		return performance;
	}

	public static Map<String, List<SingleTestCasePerformance>> compareAndReturn(
			List<EpisodicHierarchyInstance> testCases, DBN dbn, String[] targets, int memsNum,
			GraphPaintingDevice device) throws IOException {
		MemoryCreator[] strategies = new MemoryCreator[] {
				new MinimizeKLMemoryCreator(device),
				new PointSurpriseMemoryCreator(device),
				new IntervalSurpriseMemoryCreator(device) };

		return compareAndReturn(testCases, dbn, targets, strategies, memsNum, device);
	}

	public static Map<String, List<SingleTestCasePerformance>> compareAndReturn(
			List<EpisodicHierarchyInstance> testCases, DBN dbn,
			String[] targets, MemoryCreator[] strategies, int memsNum,
			GraphPaintingDevice device) throws IOException {
		final PrintStream out = System.out;

		File root = null;
		if (device != null) {
			root = device.getRootDirectory();
		}

		// init map
		Map<String, List<SingleTestCasePerformance>> strategyToResults = new HashMap<>();
		for (MemoryCreator creator : strategies) {
			strategyToResults.put(creator.getStrategyName(),
					new ArrayList<SingleTestCasePerformance>());
		}

		int i = 0;
		for (EpisodicHierarchyInstance testCase : testCases) {
			out.println("Target");
			out.println(testCase.toString());

			// compute memories
			for (MemoryCreator creator : strategies) {
				if (device != null) {
					File caseDir = new File(new File(root,
							creator.getStrategyName()), Integer.toString(i));
					caseDir.mkdirs();
					device.setRootDirectory(caseDir);
				}

				// use only observations as input for memory creation
				List<MemoryElement> mems = RetrospectiveMemoriesCreator
						.createMemories(dbn,
								testCase.getHierarchyByTime().get(0), targets,
								memsNum, creator, device);
				SingleTestCasePerformance performance = computePerformance(
						testCase, dbn, mems);
				strategyToResults.get(creator.getStrategyName()).add(
						performance);
			}

			i++;
		}
		return strategyToResults;
	}

	public static void compare(List<String> observation, DBN dbn,
			String[] targets, GraphPaintingDevice device) throws IOException {
		final int memsNum = 3;
		final PrintStream out = System.out;

		MemoryCreator[] strategies = new MemoryCreator[] {
				new MinimizeKLMemoryCreator(device),
				new PointSurpriseMemoryCreator(device),
				new IntervalSurpriseMemoryCreator(device) };

		List<List<MemoryElement>> mems = new ArrayList<List<MemoryElement>>();

		// compute memories
		for (MemoryCreator creator : strategies) {
			mems.add(RetrospectiveMemoriesCreator.createMemories(dbn,
					observation, targets, memsNum, creator, device));
		}

		for (int i = 0; i < memsNum; i++) {
			out.println("Target");
			out.println(StringUtils.join(observation, ' '));

			for (int j = 0; j < strategies.length; j++) {
				List<MemoryElement> memories = mems.get(j).subList(0, i + 1);
				dbn.resetEvidence(memories);
				dbn.getNet().updateBeliefs();
				List<String> recall = Arrays.asList(dbn.getDBNNodeMostProbableStringValues("O"));

				out.print(strategies[j].getClass().getSimpleName() + "\t"
						+ computeMatch(observation, recall));
				out.println(StringUtils.join(memories, ' '));
				out.println(StringUtils.join(recall, ' '));
			}
		}
	}

	public static double computeMatch(List a, List b) {
		int matches = 0;
		for (int i = 0; i < a.size(); i++) {
			if (a.get(i).equals(b.get(i)))
				matches++;
		}
		return matches / (double) a.size();
	}
}
